package com.hmj.hmjaiagent.agent.model;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.stream.Collectors;

@Slf4j
public class ToolCallAgent extends ReActAgent{

    private ToolCallback[] tools;

    private ChatResponse toolCallChatResponse;

    private final ToolCallingManager toolCallingManager;

    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] tools) {
        super();
        this.tools = tools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions.builder()
                //禁用Spring AI工具自主调用机制，自定义维护选项与上下文
                .withProxyToolCalls(true)
                .build();
    }

    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()) {
            return "无工具调用";
        }

        Prompt prompt = new Prompt(getMessages(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        setMessages(toolExecutionResult.conversationHistory());
        ToolResponseMessage toolResponsemessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponsemessage.getResponses().stream()
                .map(
                        response ->
                                "工具" + response.name() + "执行完成，结果：" + response.responseData())
                .collect(Collectors.joining("\n"));
        // 判断是否使用了终止工具
        boolean isTerminate = toolResponsemessage.getResponses().stream()
                .anyMatch(response ->
                        "doTerminate".equals(response.name()));
        if (isTerminate) {
            setState(AgentState.FINISHED);
        }
        log.info(results);
        return results;
    }

    @Override
    public Boolean think() {
        if (getNextStepPrompt() != null && !getNextStepPrompt().isEmpty()) {
            Message userMessage = new UserMessage(getNextStepPrompt());
            getMessages().add(userMessage);
        }

        Prompt prompt = new Prompt(getMessages(), chatOptions);

        try {
        this.toolCallChatResponse = getChatClient().prompt(prompt)
                .system(getSystemPrompt())
                .tools(tools)
                .call()
                .chatResponse();
        AssistantMessage assistantMessage = toolCallChatResponse.getResult().getOutput();
        log.info(getAgentName() + "思考：" + assistantMessage.getText());
        log.info(getAgentName() + "选择" + assistantMessage.getToolCalls().size() + "工具");
        String toolCallInfo = assistantMessage.getToolCalls().stream().map(toolCall ->
                String.format("工具名称：%s, 参数%s", toolCall.name(),
                        toolCall.arguments())
        ).collect(Collectors.joining("\n"));
        log.info(toolCallInfo);

        if (!toolCallChatResponse.hasToolCalls()){
            getMessages().add(assistantMessage);
            return false;
        }else {
            return true;
        }
        } catch (Exception e) {
            log.error(getAgentName() + "思考出错：" + e.getMessage());
            return false;
        }
    }

    @Override
    public String clean() {
        return "";
    }
}
